# 导入必要的库
import pandas as pd
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

# 配置 matplotlib 以支持中文显示及正确显示负号
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

from numpy import array
from sklearn import metrics
from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, LSTM, Bidirectional
from numpy.random import seed

# 查看 numpy 的相关属性（此处可根据实际需求决定是否保留，暂保留示例原样）
np.__all__
np.__dict__
np.__doc__
np.__file__
# np.__package__
dir(np)

# 导入股票数据获取模块 tushare
import tushare as ts

# 查看 tushare 模块包含的内容（此处可根据实际需求决定是否保留，暂保留示例原样）
dir(ts)

# Tushare 初始化相关操作

# 以下代码被注释掉了，可能是之前尝试获取数据的方式，可根据实际需求决定是否启用
# pro=ts.get_hist_data('603722')
help(ts.pro_api)

# 接口说明：
# A股日线行情接口为 daily，数据入库时间为交易日每天15点 - 16点之间，是未复权行情，停牌期间无数据。
# 调取限制：120积分每分钟内最多调取500次，每次6000条数据，单次提取量相当于23年历史数据。
# 可通过该接口获取股票行情数据，包含前后复权数据等情况，以下是接口输入输出参数说明。

# 输入参数说明
# |名称|类型|必选|描述|
# |----|----|----|----|
# |ts_code|str|N|股票代码（支持多个股票同时提取，逗号分隔）|
# |trade_date|str|N|交易日期（格式为 YYYYMMDD）|
# |start_date|str|N|开始日期(格式为 YYYYMMDD)|
# |end_date|str|N|结束日期(格式为 YYYYMMDD)|

# 注：日期都需按 YYYYMMDD 格式填写，比如 20181010。

# 输出参数说明
# |名称|类型|描述|
# |----|----|----|
# |ts_code|str|股票代码|
# |trade_date|str|交易日期|
# |open|float|开盘价|
# |high|float|最高价|
# |low|float|最低价|
# |close|float|收盘价|
# |pre_close|float|昨收价(前复权)|
# |change|float|涨跌额|
# |pct_chg|float|涨跌幅 （未复权，如果是复权请用对应通用行情接口）|
# |vol|float|成交量 （手）|
# |amount|float|成交额 （千元）|

# 初始化 pro 接口，需在 https://tushare.pro/ 注册并获取 token，这里替换为真实有效的 token
pro = ts.pro_api('572b6125c826e36619a670a9943995412179dabcfb739c5c48913f46')

# 拉取股票数据，这里以股票代码 '601857.SH' 为例，获取 20100101 至 20241225 期间的数据，并指定需要的字段
df = pro.daily(
    ts_code="601857.SH",
    start_date=20100101,
    end_date=20241225,
    fields=[
        "ts_code",
        "trade_date",
        "open",
        "high",
        "low",
        "close",
        "pre_close",
        "change",
        "pct_chg",
        "vol",
        "amount"
    ]
)
print(df)

# 存储数据到本地 CSV 文件，文件名设为 "stock601857.csv"，同时记录股票名称和代码
print(df)
df.to_csv("stock601857.csv")
stock_name = '中国石油'
stock_code = '601857'

# 对数据进行转置操作（此处可根据后续具体需求判断转置的必要性，暂保留示例原样）
df.T

# 绘制收盘价的图形，设置图形尺寸及显示图例
df["close"].plot(figsize=(16, 4), legend=True)

# 使用 pandas 读取本地存储的股票数据，将 'trade_date' 列解析为日期格式，并设为索引，同时选取需要的列
mydata = pd.read_csv("./stock601857.csv", parse_dates=["trade_date"], index_col="trade_date")[["open", "high", "low", "close"]]
# 翻转数据（可能是为了按时间顺序等需求进行调整，具体取决于数据使用场景）
mydata = mydata.iloc[::-1]
plt.plot(mydata["close"])

# 获取股价相关数据，取过去30天数据（此处索引调整可根据实际数据情况和需求进一步确认合理性）
df3 = mydata.reset_index().iloc[30:, :6]
# 去除包含空值的行，并重新从零开始编号索引
df3 = df3.dropna(how='any').reset_index(drop=True)
# 按照交易日期升序排列数据
df3 = df3.sort_values(by='trade_date', ascending=True)
print(df3.info())

# 计算均线数据，分别计算 5 日、10 日均线
df3['5'] = df3.close.rolling(5).mean()
df3['10'] = df3.close.rolling(10).mean()
df3.tail()

# 导入相关绘图模块，并配置 matplotlib 绘图样式
import matplotlib
matplotlib.style.use('ggplot')
import mplfinance as mpf
from mplfinance.original_flavor import candlestick2_ohlc
from matplotlib.ticker import FormatStrFormatter

# 创建图形和坐标轴对象，设置图形尺寸和分辨率
fig, ax = plt.subplots(1, 1, figsize=(8, 3), dpi=200)

# 绘制蜡烛图，设置颜色等参数
candlestick2_ohlc(ax,
                  opens=df3['open'].values,
                  highs=df3['high'].values,
                  lows=df3['low'].values,
                  closes=df3['close'].values,
                  width=0.5, colorup="r", colordown="g")

# 在图形上显示最高价和最低价对应的文本信息，设置字体大小
ax.text(df3.high.idxmax(), df3.high.max(), s=df3.high.max(), fontsize=8)
ax.text(df3.high.idxmin(), df3.high.min() - 2, s=df3.high.min(), fontsize=8)

# 设置坐标轴背景颜色和图形标题，设置标题字体大小
ax.set_facecolor("white")
ax.set_title(stock_name, fontsize=8)

# 绘制均线，设置透明度和图例标签
plt.plot(df3['5'].values, alpha=0.5, label='MA5')
plt.plot(df3['10'].values, alpha=0.5, label='MA10')

# 设置图例相关属性，包括背景色、边框色和字体大小
ax.legend(facecolor='white', edgecolor='white', fontsize=6)

# 修改 x 轴坐标刻度，设置刻度标签格式及旋转角度、字体大小
plt.xticks(ticks=np.arange(0, len(df3)), labels=df3.trade_date.dt.strftime('%Y-%m-%d').to_numpy())
plt.xticks(rotation=45, size=8)

# 修改 y 轴坐标刻度格式
ax.yaxis.set_major_formatter(FormatStrFormatter('%.2f'))

# 显示图形
plt.show()

# 获取股价数据（此处重新获取全部数据，可根据实际情况判断是否重复操作，暂保留示例原样）
df3 = mydata.reset_index().iloc[:, :]
# 去除包含空值的行，并重新从零开始编号索引
df3 = df3.dropna(how='any').reset_index(drop=True)
print(df3)
# 按照交易日期升序排列数据
df3 = df3.sort_values(by='trade_date', ascending=True)
print(df3.info())

# 计算均线数据，新增 30 日均线
df3['5'] = df3.close.rolling(5).mean()
df3['10'] = df3.close.rolling(10).mean()
df3['30'] = df3.close.rolling(30).mean()
df3

# 配置 matplotlib 绘图样式
matplotlib.style.use('ggplot')
from mplfinance.original_flavor import candlestick2_ohlc
from matplotlib.ticker import FormatStrFormatter

# 创建图形和坐标轴对象，设置图形尺寸和分辨率
fig, ax = plt.subplots(1, 1, figsize=(8, 3), dpi=200)

# 绘制蜡烛图，设置颜色等参数
candlestick2_ohlc(ax,
                  opens=df3['open'].values,
                  highs=df3['high'].values,
                  lows=df3['low'].values,
                  closes=df3['close'].values,
                  width=0.5, colorup="r", colordown="g")

# 在图形上显示最高价和最低价对应的文本信息，设置字体大小
ax.text(df3.high.idxmax(), df3.high.max(), s=df3.high.max(), fontsize=8)
ax.text(df3.high.idxmin(), df3.high.min() - 2, s=df3.high.min(), fontsize=8)

# 设置坐标轴背景颜色和图形标题，设置标题字体大小
ax.set_facecolor("white")
ax.set_title(stock_name + "股价走势", fontsize=8)

# 绘制均线，设置透明度和图例标签
plt.plot(df3['5'].values, alpha=0.5, label='MA5')
plt.plot(df3['10'].values, alpha=0.5, label='MA10')
plt.plot(df3['30'].values, alpha=0.5, label='MA30')

# 设置图例相关属性，包括背景色、边框色和字体大小
ax.legend(facecolor='white', edgecolor='white', fontsize=6)

# 修改 x 轴坐标刻度相关设置，包括刻度位置、标签、旋转角度、字体大小以及设置主刻度间隔等
tempXticks = np.arange(0, len(df3))
nameXticks = df3.trade_date.dt.strftime('%Y-%m-%d').to_numpy()
plt.xticks(ticks=tempXticks, labels=nameXticks)
plt.xticks(rotation=45, size=8)
x_major_locator = plt.MultipleLocator(10)
ax.xaxis.set_major_locator(x_major_locator)
ax.spines['bottom'].set_color('red')
ax.spines['left'].set_color('red')
plt.xlim(0, 100)

# 修改 y 轴坐标刻度格式
ax.yaxis.set_major_formatter(FormatStrFormatter('%.2f'))

# 显示图形
plt.show()

# 绘制开盘价和收盘价的折线图，并显示图例
data1 = mydata
data1["open"].plot(label="open")
data1["close"].plot(label="close")
plt.legend()
plt.show()

# 以下部分代码逻辑较复杂，功能可能是对数据进行整理转换，可根据实际需求进一步优化其实现方式或添加更清晰的注释说明
ddd = []
for item in data1.items():
    ddd.append(item)
ind = ddd[0][1].values
da1 = ddd[1][1].values
index1 = []
data2 = []
for item in ind:
    index1.append(item)
index1.reverse()
for item in da1:
    data2.append(item)
data2.reverse()
data1 = {'trade_date': index1, 'close': data2}
stockdata = pd.DataFrame(data1)

# 数据检查操作（可根据具体需求确认是否需要输出此部分内容，暂保留示例原样）
df3.iloc[:, 4:5].values[:, :]

# 将交易日期转换为数值格式，并添加为新列 'date'（具体后续用途可根据代码逻辑进一步完善注释）
from matplotlib.dates import date2num
date2num(df3.trade_date.values)
df3['date'] = date2num(df3.trade_date.values)
# df3=df3.sort_values(by="date",ascending=True)
# df3=df3[['date','open', 'high', 'low', 'close']]

ind = np.arange(0, 2677)
dataf1 = df3[['date', 'open', 'high', 'low', 'close']]

# 设置索引（需确认此处设置索引是否达到预期效果，暂保留示例原样）
dataf1.set_index("date")

# 深度学习拟合股票相关操作

# 为保证结果可复现，设置随机种子
seed(1)
tf.random.set_seed(1)

# 设置相关参数，包括时间戳数量和训练轮数，以及选择模型的类型（可考虑添加更详细的注释说明这些参数的含义及对模型的影响）
n_timestamp = 5
n_epochs = 30
model_type = 1

# 展示数据（此处可根据具体分析或调试需求判断是否需要展示，暂保留示例原样）
dataf1

# 拟合开始，数据准备阶段
# 提取训练集和测试集数据，训练集取前 2048 行的收盘价数据，测试集取后 300 行左右的收盘价数据（索引计算可进一步确认准确性）
training_set = dataf1.iloc[0:2048, 4:5].to_numpy()
test_set = dataf1.iloc[3626 - 300:, 4:5].to_numpy()
plt.plot(training_set)

# 数据归一化操作，将数据归一化到 0 到 1 的范围，分别对训练集和测试集进行处理
sc = MinMaxScaler(feature_range=(0, 1))
training_set_scaled = sc.fit_transform(training_set)
testing_set_scaled = sc.transform(test_set)

# 数据分割函数，按照时间戳数量划分数据为输入特征 X 和目标值 y
def data_split(sequence, n_timestamp):
    X = []
    y = []
    for i in range(len(sequence)):
        end_ix = i + n_timestamp
        if end_ix > len(sequence) - 1:
            break
        seq_x, seq_y = sequence[i:end_ix], sequence[end_ix]
        X.append(seq_x)
        y.append(seq_y)
    return array(X), array(y)

X_train, y_train = data_split(training_set_scaled, n_timestamp)

# 训练数据重整，调整数据形状以适配模型输入要求，对训练集和测试集分别进行操作
X_train = X_train.reshape(X_train.shape[0], X_train.shape[1], 1)
X_test, y_test = data_split(testing_set_scaled, n_timestamp)
X_test = X_test.reshape(X_test.shape[0], X_test.shape[1], 1)

# 构建模型，根据选择的模型类型（1: 单层 LSTM，2: 多层 LSTM，3: 双向 LSTM）构建相应的深度学习模型，并输出模型结构信息
model_type = 2
if model_type == 1:
    # 单层 LSTM 模型构建
    model = Sequential()
    model.add(LSTM(units=50, activation='relu',
                   input_shape=(X_train.shape[1], 1)))
    model.add(Dense(units=1))
if model_type == 2:
    # 多层 LSTM 模型构建
    model = Sequential()
    model.add(LSTM(units=50, activation='relu', return_sequences=True,
                   input_shape=(X_train.shape[1], 1)))
    model.add(LSTM(units=50, activation='relu'))
    model.add(Dense(1))
if model_type == 3:
    # 双向 LSTM 模型构建
    model = Sequential()
    model.add(Bidirectional(LSTM(50, activation='relu'),
                            input_shape=(X_train.shape[1], 1)))
    model.add(Dense(1))
model.summary()

# 模型编译，设置优化器为 Adam，学习率# 模型编译，设置优化器为Adam，学习率为0.001，损失函数选用均方误差
model.compile(optimizer=tf.keras.optimizers.Adam(0.001),
              loss='mean_squared_error')

# 训练模型，指定批次大小、训练轮数、验证数据以及验证频率，并输出模型结构信息（可考虑添加训练过程的进度显示等优化）
history = model.fit(X_train, y_train,
                    batch_size=64,
                    epochs=n_epochs,
                    validation_data=(X_test, y_test),
                    validation_freq=1)
model.summary()

# 输出训练结果，绘制训练损失和验证损失的变化曲线，方便观察模型训练情况
plt.plot(history.history['loss'], label='Training Loss')
plt.plot(history.history['val_loss'], label='Validation Loss')
plt.legend()
plt.show()

# 模型预测阶段，使用训练好的模型对测试集数据进行预测，并将预测结果和真实数据进行反归一化操作，还原到原始数据范围
predicted_stock_price = model.predict(X_test)
predicted_stock_price = sc.inverse_transform(predicted_stock_price)
real_stock_price = sc.inverse_transform(y_test)

# 画出真实数据和预测数据的对比曲线，直观展示模型预测效果，设置图形标题、坐标轴标签以及图例等信息
plt.plot(real_stock_price, color='red', label='Stock Price')
plt.plot(predicted_stock_price, color='blue', label='Predicted Stock Price')
plt.title(stock_name)
plt.xlabel('Time')
plt.ylabel('Stock Price')
plt.legend()
plt.show()

# 模型校核阶段，计算均方误差（MSE）、均方根误差（RMSE）、平均绝对误差（MAE）以及决定系数（R2），评估模型性能
MSE = metrics.mean_squared_error(predicted_stock_price, real_stock_price)
RMSE = metrics.mean_squared_error(predicted_stock_price, real_stock_price) ** 0.5
MAE = metrics.mean_absolute_error(predicted_stock_price, real_stock_price)
R2 = metrics.r2_score(predicted_stock_price, real_stock_price)

print('均方误差: %.5f' % MSE)
print('均方根误差: %.5f' % RMSE)
print('平均绝对误差: %.5f' % MAE)
print('R2: %.5f' % R2)